{-
 - Copyright (C) 2019  Koz Ross <koz.ross@retro-freedom.nz>
 -
 - This program is free software: you can redistribute it and/or modify
 - it under the terms of the GNU General Public License as published by
 - the Free Software Foundation, either version 3 of the License, or
 - (at your option) any later version.
 -
 - This program is distributed in the hope that it will be useful,
 - but WITHOUT ANY WARRANTY; without even the implied warranty of
 - MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 - GNU General Public License for more details.
 -
 - You should have received a copy of the GNU General Public License
 - along with this program.  If not, see <http://www.gnu.org/licenses/>.
 -}

{-# OPTIONS_GHC -fplugin GHC.TypeLits.Normalise #-}
{-# OPTIONS_GHC -fplugin GHC.TypeLits.KnownNat.Solver #-}

{-# LANGUAGE Trustworthy #-}
{-# LANGUAGE ConstrainedClassMethods #-}
{-# LANGUAGE TypeInType #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE TypeOperators #-}
{-# LANGUAGE TypeApplications #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE DefaultSignatures #-}
{-# LANGUAGE UndecidableInstances #-}
{-# LANGUAGE TemplateHaskell #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE CPP #-}

#if MIN_VERSION_base(4,12,0)
{-# LANGUAGE NoStarIsType #-}
#endif

-- | 
-- Module:        Data.Finitary
-- Description:   A type class witnessing that a type has finite cardinality.
-- Copyright:     (C) Koz Ross, 2019
-- License:       GPL version 3.0 or later
-- Maintainer:    koz.ross@retro-freedom.nz
-- Stability:     Experimental
-- Portability:   GHC only
--
-- This package provides the 'Finitary' type class, as well as a range of useful
-- \'base\' instances for commonly-used finitary types. 
--
-- For your own types, there are three possible ways to define an instance of
-- 'Finitary':
--
-- __Via 'Generic'__
--
-- If your data type implements 'Generic' (and is finitary), you can
-- automatically derive your instance:
--
-- > {-# LANGUAGE DeriveAnyClass #-}
-- > {-# LANGUAGE DeriveGeneric #-}
-- >
-- > import GHC.Generics
-- > import Data.Word
-- >
-- > data Foo = Bar | Baz (Word8, Word8) | Quux Word16
-- >    deriving (Eq, Generic, Finitary)
--
-- This is the easiest method, and also the safest, as GHC will automatically
-- determine the cardinality of @Foo@, as well as defining law-abiding methods.
-- It may be somewhat slower than a \'hand-rolled\' method in some cases.
--
-- __By defining only 'Cardinality', 'fromFinite' and 'toFinite'__
--
-- If you want a manually-defined instance, but don't wish to define every
-- method, only 'fromFinite' and 'toFinite' are needed, along with
-- 'Cardinality'. 'Cardinality' in particular must be defined with care, as
-- otherwise, you may end up with inconstructable values or indexes that don't 
-- correspond to anything.
--
-- __By defining everything__
--
-- For maximum control, you can define all the methods. Ensure you follow all
-- the laws!
--
module Data.Finitary (
  Finitary(..),
  inhabitants, inhabitantsFrom, inhabitantsTo, inhabitantsFromTo
) where

import Data.List.NonEmpty (NonEmpty(..))
import Data.Bifunctor (bimap, first)
import Numeric.Natural (Natural)
import Data.Semigroup (Max, Min, Sum, Product, Dual, Last, First, Any, All)
import Data.Functor.Identity (Identity)
import Data.Int (Int8, Int16, Int32, Int64)
import Data.Word (Word8, Word16, Word32, Word64)
import Data.Proxy (Proxy(..))
import Data.Void (Void)
import Data.Bool (bool)
import CoercibleUtils (op)
import GHC.Generics (Generic, Rep, U1(..), K1(..), V1, (:+:)(..), (:*:)(..), M1(..), from, to)
import Control.Applicative (Alternative(..), Const)
import Data.Kind (Type)
import GHC.TypeNats
import Data.Finite (Finite, finites, separateSum, separateProduct, combineProduct, weakenN, shiftN)
import Data.Ord (Down(..))
import Control.Monad.Primitive (PrimMonad(..))
import Control.Monad (forM_, join)
import GHC.TypeLits.Compare (isLE)
import Data.Type.Equality ((:~:)(..))
import Control.Monad.ST (ST, runST)
import Foreign.Storable (Storable)

import qualified Data.List.NonEmpty as NE
import qualified Data.Bit as B
import qualified Data.Bit.ThreadSafe as BTS
import qualified Data.Vector.Sized as VS
import qualified Data.Vector.Generic as VG
import qualified Data.Vector.Mutable.Sized as VMS
import qualified Data.Vector.Generic.Sized as VGS
import qualified Data.Vector.Generic.Mutable as VGM
import qualified Data.Vector.Generic.Mutable.Sized as VGMS
import qualified Data.Vector.Unboxed.Sized as VUS
import qualified Data.Vector.Unboxed.Mutable.Sized as VUMS
import qualified Data.Vector.Storable.Sized as VSS
import qualified Data.Vector.Storable.Mutable.Sized as VSMS

import Data.Finitary.TH

-- | Witnesses an isomorphism between @a@ and some @(KnownNat n) => Finite n@.
-- Effectively, a lawful instance of this shows that @a@ has exactly @n@
-- (non-@_|_@) inhabitants, and that we have a bijection with 'fromFinite' and
-- 'toFinite' as each \'direction\'.
--
-- For any type @a@ with an instance of @Finitary@, for every non-@_|_@ @x :: a@, we have
-- a unique /index/ @i :: Finite n@. We will also refer to any such @x@ as an
-- /inhabitant/ of @a@. We can convert inhabitants to indexes using @toFinite@,
-- and also convert indexes to inhabitants with @fromFinite@.
--
-- __Laws__
--
-- The main laws state that 'fromFinite' should be a bijection, with 'toFinite' as
-- its inverse, and 'Cardinality' must be a truthful representation of the
-- cardinality of the type. Thus:
--
--    * \[\texttt{fromFinite} \circ \texttt{toFinite} = \texttt{toFinite} \circ
-- \texttt{fromFinite} = \texttt{id}\]
--    * \[\forall x, y :: \texttt{Finite} \; (\texttt{Cardinality} \; a) \; \texttt{fromFinite} \; x = \texttt{fromFinite} \; y
-- \rightarrow x = y\]
--    * \[\forall x :: \texttt{Finite} \; (\texttt{Cardinality} \; a) \; \exists y :: a \mid \texttt{fromFinite} \; x
-- = y\]
--
-- Furthermore, 'fromFinite' should be _order-preserving_. Namely, if @a@ is an
-- instance of @Ord@, we must have:
--
--    * \[\forall i, j :: \texttt{Finite} \; (\texttt{Cardinality} \; a) \;
--    \texttt{fromFinite} \; i \leq \texttt{fromFinite} \; j \rightarrow i \leq j \]
--
-- Lastly, if you define any of the other methods, these laws must hold:
--
--    * \[ a \neq \emptyset \rightarrow \texttt{start} = \texttt{fromFinite} \; \texttt{minBound} \]
--    * \[ a \neq \emptyset \rightarrow \texttt{end} = \texttt{fromFinite} \; \texttt{maxBound} \]
--    * \[ \forall x :: a \; \texttt{end} \neq x \rightarrow \texttt{next} \; x =
-- (\texttt{fromFinite} \circ + 1 \circ \texttt{toFinite}) \; x \]
--    * \[ \forall x :: a \; \texttt{start} \neq x \rightarrow \texttt{previous} \; x =
-- (\texttt{fromFinite} \circ - 1 \circ \texttt{toFinite}) \; x \]
-- 
-- Together with the fact that @Finite n@ is well-ordered whenever @KnownNat n@
-- holds, a law-abiding @Finitary@ instance for a type @a@ defines a constructive
-- [well-order](https://en.wikipedia.org/wiki/Well-order), witnessed by
-- 'toFinite' and 'fromFinite', which agrees with the @Ord@ instance for @a@, if
-- any.
-- 
-- We /strongly/ suggest that @fromFinite@ and @toFinite@ should have
-- time complexity \(\Theta(1)\), or, if that's not possible, \(O(\texttt{Cardinality} \; a)\). 
-- The latter is the case for instances generated using
-- @Generics@-based derivation, but not for \'basic\' types; thus, these
-- functions for your derived types will only be as slow as their \'structure\',
-- rather than their \'contents\', provided the contents are of these \'basic\'
-- types. 
class (Eq a, KnownNat (Cardinality a)) => Finitary (a :: Type) where
  -- | How many (non-@_|_@) inhabitants @a@ has, as a typelevel natural number. 
  type Cardinality a :: Nat
  type Cardinality a = GCardinality (Rep a)
  -- | Converts an index into its corresponding inhabitant.
  fromFinite :: Finite (Cardinality a) -> a
  default fromFinite :: (Generic a, GFinitary (Rep a), Cardinality a ~ GCardinality (Rep a)) => Finite (Cardinality a) -> a
  fromFinite = to . gFromFinite
  -- | Converts an inhabitant to its corresponding index.
  toFinite :: a -> Finite (Cardinality a)
  default toFinite :: (Generic a, GFinitary (Rep a), Cardinality a ~ GCardinality (Rep a)) => a -> Finite (Cardinality a)
  toFinite = gToFinite . from
  -- | The first inhabitant, by index, assuming @a@ has any inhabitants.
  start :: (1 <= Cardinality a) => a
  start = fromFinite minBound
  -- | The last inhabitant, by index, assuming @a@ has any inhabitants.
  end :: (1 <= Cardinality a) => a
  end = fromFinite maxBound
  -- | @previous x@ gives the inhabitant whose index precedes the index of @x@,
  -- or 'empty' if no such index exists.
  previous :: (Alternative f) => a -> f a
  previous = fmap fromFinite . guarded (/= maxBound) . dec . toFinite
  -- | @next x@ gives the inhabitant whose index follows the index of @x@, or
  -- 'empty' if no such index exists.
  next :: (Alternative f) => a -> f a
  next = fmap fromFinite . guarded (/= minBound) . inc . toFinite

class (KnownNat (GCardinality a)) => GFinitary (a :: Type -> Type) where
  type GCardinality a :: Nat
  gFromFinite :: Finite (GCardinality a) -> a x
  gToFinite :: a x -> Finite (GCardinality a) 

instance GFinitary V1 where
  type GCardinality V1 = 0
  {-# INLINE gFromFinite #-}
  gFromFinite = const undefined
  {-# INLINE gToFinite #-}
  gToFinite = const undefined 

instance GFinitary U1 where
  type GCardinality U1 = 1
  {-# INLINE gFromFinite #-}
  gFromFinite = const U1
  {-# INLINE gToFinite #-}
  gToFinite = const 0

instance (Finitary a) => GFinitary (K1 _1 a) where
  type GCardinality (K1 _1 a) = Cardinality a
  {-# INLINE gFromFinite #-}
  gFromFinite = K1 . fromFinite
  {-# INLINE gToFinite #-}
  gToFinite = toFinite . op K1

instance (GFinitary a, GFinitary b) => GFinitary (a :+: b) where
  type GCardinality (a :+: b) = GCardinality a + GCardinality b
  {-# INLINE gFromFinite #-}
  gFromFinite = either (L1 . gFromFinite) (R1 . gFromFinite) . separateSum
  {-# INLINE gToFinite #-}
  gToFinite (L1 x) = weakenN . gToFinite $ x
  gToFinite (R1 x) = shiftN . gToFinite $ x 

instance (GFinitary a, GFinitary b) => GFinitary (a :*: b) where
  type GCardinality (a :*: b) = GCardinality a * GCardinality b
  {-# INLINE gFromFinite #-}
  gFromFinite i = let (x, y) = separateProduct' i in
                    gFromFinite x :*: gFromFinite y
  {-# INLINE gToFinite #-}
  gToFinite (x :*: y) = combineProduct' @(GCardinality a) @(GCardinality b) (weakenN . gToFinite $ x, weakenN . gToFinite $ y)

instance (GFinitary a) => GFinitary (M1 _x _y a) where
  type GCardinality (M1 _x _y a) = GCardinality a
  {-# INLINE gFromFinite #-}
  gFromFinite = M1 . gFromFinite
  {-# INLINE gToFinite #-}
  gToFinite = gToFinite . op M1

-- * Instances
-- Basic types

instance Finitary Void

instance Finitary ()

instance Finitary (Proxy a)

instance Finitary Bool

instance Finitary Any

instance Finitary All

instance Finitary B.Bit where
  type Cardinality B.Bit = 2
  {-# INLINE fromFinite #-}
  fromFinite = B.Bit . toEnum . fromEnum
  {-# INLINE toFinite #-}
  toFinite = toEnum . fromEnum . op B.Bit
  {-# INLINE start #-}
  start = minBound
  {-# INLINE end #-}
  end = maxBound
  {-# INLINE next #-}
  next = fmap succ . guarded (== minBound)
  {-# INLINE previous #-}
  previous = fmap pred . guarded (== maxBound)

instance Finitary BTS.Bit where
  type Cardinality BTS.Bit = 2
  {-# INLINE fromFinite #-}
  fromFinite = BTS.Bit . toEnum . fromEnum
  {-# INLINE toFinite #-}
  toFinite = toEnum . fromEnum . op BTS.Bit
  {-# INLINE start #-}
  start = minBound
  {-# INLINE end #-}
  end = maxBound
  {-# INLINE next #-}
  next = fmap succ . guarded (== minBound)
  {-# INLINE previous #-}
  previous = fmap pred . guarded (== maxBound)

instance Finitary Ordering

-- | 'Char' has one inhabitant per Unicode code point.
instance Finitary Char where
  type Cardinality Char = $(charCardinality)
  {-# INLINE fromFinite #-}
  fromFinite = toEnum . fromEnum
  {-# INLINE toFinite #-}
  toFinite = toEnum . fromEnum
  {-# INLINE start #-}
  start = minBound
  {-# INLINE end #-}
  end = maxBound
  {-# INLINE next #-}
  next = fmap succ . guarded (/= maxBound)
  {-# INLINE previous #-}
  previous = fmap pred . guarded (/= minBound)

instance Finitary Word8 where
  type Cardinality Word8 = $(cardinalityOf @Word8)
  {-# INLINE fromFinite #-}
  fromFinite = toEnum . fromEnum
  {-# INLINE toFinite #-}
  toFinite = toEnum . fromEnum
  {-# INLINE start #-}
  start = minBound
  {-# INLINE end #-}
  end = maxBound
  {-# INLINE next #-}
  next = fmap succ . guarded (/= maxBound)
  {-# INLINE previous #-}
  previous = fmap pred . guarded (/= minBound)

instance Finitary Word16 where
  type Cardinality Word16 = $(cardinalityOf @Word16)
  {-# INLINE fromFinite #-}
  fromFinite = toEnum . fromEnum
  {-# INLINE toFinite #-}
  toFinite = toEnum . fromEnum
  {-# INLINE start #-}
  start = minBound
  {-# INLINE end #-}
  end = maxBound
  {-# INLINE next #-}
  next = fmap succ . guarded (/= maxBound)
  {-# INLINE previous #-}
  previous = fmap pred . guarded (/= minBound)

instance Finitary Word32 where
  type Cardinality Word32 = $(cardinalityOf @Word32)
  {-# INLINE fromFinite #-}
  fromFinite = fromIntegral
  {-# INLINE toFinite #-}
  toFinite = fromIntegral
  {-# INLINE start #-}
  start = minBound
  {-# INLINE end #-}
  end = maxBound
  {-# INLINE next #-}
  next = guarded (== minBound) . inc
  {-# INLINE previous #-}
  previous = guarded (== maxBound) . dec

instance Finitary Word64 where
  type Cardinality Word64 = $(cardinalityOf @Word64)
  {-# INLINE fromFinite #-}
  fromFinite = fromIntegral
  {-# INLINE toFinite #-}
  toFinite = fromIntegral
  {-# INLINE start #-}
  start = minBound
  {-# INLINE end #-}
  end = maxBound
  {-# INLINE next #-}
  next = guarded (== minBound) . inc
  {-# INLINE previous #-}
  previous = guarded (== maxBound) . dec

instance Finitary Int8 where
  type Cardinality Int8 = $(cardinalityOf @Int8)
  {-# INLINE fromFinite #-}  
  fromFinite = fromIntegral . subtract 128 . fromIntegral @_ @Int16
  {-# INLINE toFinite #-}
  toFinite = fromIntegral . (+ 128) . fromIntegral @_ @Int16
  {-# INLINE start #-}
  start = minBound
  {-# INLINE end #-}
  end = maxBound
  {-# INLINE next #-}
  next = fmap succ . guarded (/= maxBound)
  {-# INLINE previous #-}
  previous = fmap pred . guarded (/= minBound)

instance Finitary Int16 where
  type Cardinality Int16 = $(cardinalityOf @Int16)
  {-# INLINE fromFinite #-}  
  fromFinite = fromIntegral . subtract 32768 . fromIntegral @_ @Int32
  {-# INLINE toFinite #-}
  toFinite = fromIntegral . (+ 32768) . fromIntegral @_ @Int32
  {-# INLINE start #-}
  start = minBound
  {-# INLINE end #-}
  end = maxBound
  {-# INLINE next #-}
  next = fmap succ . guarded (/= maxBound)
  {-# INLINE previous #-}
  previous = fmap pred . guarded (/= minBound)

instance Finitary Int32 where
  type Cardinality Int32 = $(cardinalityOf @Int32)
  {-# INLINE fromFinite #-}
  fromFinite = fromIntegral @_ @Int32 . subtract $(adjustmentOf @Int32) . fromIntegral @_ @Integer
  {-# INLINE toFinite #-}
  toFinite = fromIntegral . (+ $(adjustmentOf @Int32)) . fromIntegral @_ @Integer . fromEnum
  {-# INLINE start #-}
  start = minBound
  {-# INLINE end #-}
  end = maxBound
  {-# INLINE next #-}
  next = guarded (== minBound) . inc
  {-# INLINE previous #-}
  previous = guarded (== maxBound) . dec

instance Finitary Int64 where
  type Cardinality Int64 = $(cardinalityOf @Int64)
  {-# INLINE fromFinite #-}
  fromFinite = fromIntegral @_ @Int64 . subtract $(adjustmentOf @Int64) . fromIntegral @_ @Integer
  {-# INLINE toFinite #-}
  toFinite = fromIntegral . (+ $(adjustmentOf @Int64)) . fromIntegral @_ @Integer . fromEnum
  {-# INLINE start #-}
  start = minBound
  {-# INLINE end #-}
  end = maxBound
  {-# INLINE next #-}
  next = guarded (== minBound) . inc
  {-# INLINE previous #-}
  previous = guarded (== maxBound) . dec

-- Variable-width instances

-- | 'Int' has a finite number of inhabitants, varying by platform. This
-- instance will determine this when the library is built.
instance Finitary Int where
  type Cardinality Int = $(cardinalityOf @Int)
  {-# INLINE fromFinite #-}
  fromFinite = fromIntegral @_ @Int . subtract $(adjustmentOf @Int) . fromIntegral @_ @Integer
  {-# INLINE toFinite #-}
  toFinite = fromIntegral . (+ $(adjustmentOf @Int)) . fromIntegral @_ @Integer . fromEnum
  {-# INLINE start #-}
  start = minBound
  {-# INLINE end #-}
  end = maxBound
  {-# INLINE next #-}
  next = guarded (== minBound) . inc
  {-# INLINE previous #-}
  previous = guarded (== maxBound) . dec

-- | 'Word' has a finite number of inhabitants, varying by platform. This
-- instance will determine this when the library is built.
instance Finitary Word where
  type Cardinality Word = $(cardinalityOf @Word)
  {-# INLINE fromFinite #-}
  fromFinite = fromIntegral
  {-# INLINE toFinite #-}
  toFinite = fromIntegral
  {-# INLINE start #-}
  start = minBound
  {-# INLINE end #-}
  end = maxBound
  {-# INLINE next #-}
  next = guarded (== minBound) . inc
  {-# INLINE previous #-}
  previous = guarded (== maxBound) . dec

-- | Since any type is isomorphic to itself, it follows that a \'valid\' @Finite
-- n@ (meaning that @n@ is a 'KnownNat') has finite cardinality.
instance (KnownNat n) => Finitary (Finite n) where
  type Cardinality (Finite n) = n
  {-# INLINE fromFinite #-}
  fromFinite = id
  {-# INLINE toFinite #-}
  toFinite = id
  {-# INLINE start #-}
  start = minBound
  {-# INLINE end #-}
  end = maxBound
  {-# INLINE next #-}
  next = guarded (== minBound) . inc
  {-# INLINE previous #-}
  previous = guarded (== maxBound) . dec

-- | @Maybe a@ introduces one additional inhabitant (namely, 'Nothing') to @a@.
instance (Finitary a) => Finitary (Maybe a)

-- | The sum of two finite types will also be finite, with a cardinality equal
-- to the sum of their cardinalities.
instance (Finitary a, Finitary b) => Finitary (Either a b)

-- | The product of two finite types will also be finite, with a cardinality
-- equal to the product of their cardinalities.
instance (Finitary a, Finitary b) => Finitary (a, b)

instance (Finitary a, Finitary b, Finitary c) => Finitary (a, b, c)

instance (Finitary a, Finitary b, Finitary c, Finitary d) => Finitary (a, b, c, d)

instance (Finitary a, Finitary b, Finitary c, Finitary d, Finitary e) => Finitary (a, b, c, d, e)

instance (Finitary a, Finitary b, Finitary c, Finitary d, Finitary e, Finitary f) => Finitary (a, b, c, d, e, f)

instance (Finitary a) => Finitary (Const a b)

-- | For any @newtype@-esque thing over a type with a @Finitary@ instance, we
-- can just \'inherit\' the behaviour of @a@.
instance (Finitary a) => Finitary (Sum a)

instance (Finitary a) => Finitary (Product a)

instance (Finitary a) => Finitary (Dual a)

instance (Finitary a) => Finitary (Last a)

instance (Finitary a) => Finitary (First a)

instance (Finitary a) => Finitary (Identity a)

instance (Finitary a) => Finitary (Max a)

instance (Finitary a) => Finitary (Min a)

-- | Despite the @newtype@-esque nature of @Down@, due to the requirement that
-- 'fromFinite' is order-preserving, the instance for @Down a@ reverses the
-- indexing.
instance (Finitary a) => Finitary (Down a) where
  type Cardinality (Down a) = Cardinality a
  {-# INLINE fromFinite #-}
  fromFinite = Down . fromFinite . opp
    where opp = fromIntegral @_ @(Finite (Cardinality a)) . (`mod` n) . (* (n - 1)) . (+ 1) . fromIntegral @_ @Natural 
          n = natVal @(Cardinality a) Proxy
  {-# INLINE toFinite #-}
  toFinite = fromIntegral @_ @(Finite (Cardinality a)) . (`mod` n) . (* (n - 1)) . (+ 1) . fromIntegral @_ @Natural . toFinite . op Down
    where n = natVal @(Cardinality a) Proxy 

-- | A fixed-length vector over a type @a@ with an instance of @Finitary@ can be
-- thought of as a fixed-length word over an alphabet of size @Cardinality a@.
-- Since there are only finitely-many of these, we can index them in lex order,
-- with the ordering determined by the @Finitary a@ instance (thus, the
-- \'first\' such @Vector@ is the one where each element is @start :: a@, and
-- the \'last\' is the one where each element is @end :: a@).
instance (Finitary a, KnownNat n) => Finitary (VS.Vector n a) where
  type Cardinality (VS.Vector n a) = Cardinality a ^ n
  {-# INLINE fromFinite #-}
  fromFinite i = runST (go i)
    where go :: Finite (Cardinality (VS.Vector n a)) -> ST s (VS.Vector n a)
          go ix = do v <- VMS.new
                     unroll v ix
                     VS.unsafeFreeze v
  {-# INLINE toFinite #-}
  toFinite = roll

instance (Finitary a, VUMS.Unbox a, KnownNat n) => Finitary (VUS.Vector n a) where
  type Cardinality (VUS.Vector n a) = Cardinality a ^ n
  {-# INLINE fromFinite #-}
  fromFinite i = runST (go i)
    where go :: Finite (Cardinality (VUS.Vector n a)) -> ST s (VUS.Vector n a)
          go ix = do v <- VUMS.new
                     unroll v ix
                     VUS.unsafeFreeze v
  {-# INLINE toFinite #-}
  toFinite = roll

instance (Finitary a, Storable a, KnownNat n) => Finitary (VSS.Vector n a) where
  type Cardinality (VSS.Vector n a) = Cardinality a ^ n
  {-# INLINE fromFinite #-}
  fromFinite i = runST (go i)
    where go :: Finite (Cardinality (VSS.Vector n a)) -> ST s (VSS.Vector n a)
          go ix = do v <- VSMS.new
                     unroll v ix
                     VSS.unsafeFreeze v
  {-# INLINE toFinite #-}
  toFinite = roll

-- * Enumeration helpers

-- | Produce every inhabitant of @a@, in ascending order of indexes.
-- If you want descending order, use @Down a@ instead.
{-# INLINE inhabitants #-}
inhabitants :: forall (a :: Type) . (Finitary a) => [a]
inhabitants = fromFinite <$> finites

-- | Produce every inhabitant of @a@, starting with the argument, in ascending
-- order of indexes.
-- If you want descending order, use @Down a@ instead.
{-# INLINE inhabitantsFrom #-}
inhabitantsFrom :: forall (a :: Type) . (Finitary a) => a -> NonEmpty a
inhabitantsFrom x = x :| concatMap @Maybe (fmap fromFinite . enumFrom . toFinite) (next x) 

-- | Produce every inhabitant of @a@, up to and including the argument, in
-- ascending order of indexes.
-- If you want descending order, use @Down a@ instead.
{-# INLINE inhabitantsTo #-}
inhabitantsTo :: forall (a :: Type) . (Finitary a) => a -> NonEmpty a
inhabitantsTo x = NE.fromList (fromFinite <$> [0 .. toFinite x]) 

-- | Produce every inhabitant of @a@, starting with the first argument, up to
-- the second argument, in ascending order of indexes. @inhabitantsFromTo x y@
-- will produce the empty list if @toFinite x > toFinite y@.
-- If you want descending order, use @Down a@ instead.
{-# INLINE inhabitantsFromTo #-}
inhabitantsFromTo :: forall (a :: Type) . (Finitary a) => a -> a -> [a]
inhabitantsFromTo lo hi = fromFinite <$> [toFinite lo .. toFinite hi]

-- Helpers

{-# INLINE combineProduct' #-}
combineProduct' :: forall n m . (KnownNat n, KnownNat m) => (Finite n, Finite m) -> Finite (n * m)
combineProduct' = fromIntegral . uncurry (+) . first ((natVal $ Proxy @m) *) . bimap @_ @_ @Natural @_ @Natural fromIntegral fromIntegral

{-# INLINE separateProduct' #-}
separateProduct' :: forall n m . (KnownNat n, KnownNat m) => Finite (n * m) -> (Finite n, Finite m)
separateProduct' = bimap (fromIntegral . (\x -> fromIntegral x `div` natVal @m Proxy)) (fromIntegral . (\x -> fromIntegral x `mod` natVal @m Proxy)) . join (,)

unroll :: forall a m v n . (Finitary a, PrimMonad m, KnownNat n, VGM.MVector v a) => VGMS.MVector v n (PrimState m) a -> Finite (Cardinality a ^ n) -> m ()
unroll v acc = forM_ @_ @_ @_ @() (isLE (Proxy @1) (Proxy @n)) 
                               (\Refl -> do let (d, r) = separateProduct @(Cardinality a ^ (n -1)) @(Cardinality a) acc
                                            let x = fromFinite r
                                            VGMS.write v 0 x
                                            unroll (VGMS.tail v) d)

roll :: forall a v n . (Finitary a, VG.Vector v a, KnownNat n) => VGS.Vector v n a -> Finite (Cardinality a ^ n)
roll v = case isLE (Proxy @1) (Proxy @n) of
          Nothing -> 0
          Just Refl -> let (h, t) = (VGS.head v, VGS.tail v) in
                          combineProduct (roll t, toFinite h) 

{-# INLINE inc #-}
inc :: (Num a) => a -> a
inc = (+ 1)

{-# INLINE dec #-}
dec :: (Num a) => a -> a
dec = subtract 1

{-# INLINE guarded #-}
guarded :: forall (a :: Type) (f :: Type -> Type) . (Alternative f) => (a -> Bool) -> a -> f a
guarded p x = bool empty (pure x) (p x)
